import json
import re
from rag_chain import RAGChain
from text_processor import TextVectorizer
import requests
import time

# 目标电话号码和公司
target_phone = "15332321312"
expected_company = "深圳创新科技有限公司"
# 测试文件名
test_file_path = "test_new_phone.txt"

# 1. 检查测试文件是否存在并读取内容
print("=== 检查测试文件 ===")
try:
    with open(test_file_path, 'r', encoding='utf-8') as f:
        file_content = f.read()
        print(f"测试文件内容: {file_content[:100]}...")
        print(f"文件中包含目标电话号码: {target_phone in file_content}")
        print(f"文件中包含目标公司: {expected_company in file_content}")
except Exception as e:
    print(f"读取测试文件失败: {e}")

# 2. 初始化向量器和RAG链进行本地测试
print("\n=== 本地向量器测试 ===")
vectorizer = TextVectorizer(use_tfidf=True)
rag_chain = RAGChain(deepseek_api_key="sk-a6d2da9cbeaf4cf59512a993de770cbd", vectorizer=vectorizer)

# 3. 手动添加测试数据到向量器
vectorizer.vectorize_and_store(file_content, metadata={"source": test_file_path})
print("已手动添加测试数据到向量器")

# 4. 检查词汇表
print("\n=== 词汇表检查 ===")
print(f"词汇表大小: {len(vectorizer.vocab)}")
print(f"电话号码 '{target_phone}' 在词汇表中: {target_phone in vectorizer.vocab}")
print(f"公司名称 '{expected_company}' 在词汇表中: {expected_company in vectorizer.vocab}")

# 5. 直接使用向量器查询
print("\n=== 向量器直接查询 ===")
vec_results = vectorizer.query(target_phone, n_results=3, similarity_metric="cosine")
print(f"查询结果: {json.dumps(vec_results, ensure_ascii=False, indent=2)}")

# 6. 使用RAG链查询
print("\n=== RAG链查询 ===")
rag_results = rag_chain.run(target_phone, n_results=3, similarity_metric="cosine")
print(f"RAG链结果: {json.dumps(rag_results, ensure_ascii=False, indent=2)}")

# 7. 验证本地测试结果
print("\n=== 本地测试结果验证 ===")
if "answer" in rag_results:
    print(f"AI回答: {rag_results['answer']}")
    if expected_company in rag_results['answer'] and target_phone in rag_results['answer']:
        print(f"成功! AI回答正确识别了{target_phone}属于{expected_company}")
    else:
        print(f"AI回答未能识别关联信息")
else:
    print("RAG链结果格式不正确")

# 8. 测试API上传和检索流程
print("\n=== API流程测试 ===")

# 获取JWT令牌
def get_token():
    url = "http://localhost:8001/token"
    data = {
        "username": "admin",
        "password": "admin123"
    }
    response = requests.post(url, data=data)
    if response.status_code == 200:
        return response.json()["access_token"]
    else:
        print(f"获取令牌失败: {response.status_code} - {response.text}")
        return None

# 上传测试文件
def upload_test_file(token):
    url = "http://localhost:8001/upload-document"
    headers = {
        "Authorization": f"Bearer {token}"
    }
    with open(test_file_path, "rb") as f:
        files = {
            "file": (test_file_path, f, "text/plain")
        }
        response = requests.post(url, headers=headers, files=files)
        if response.status_code == 200:
            print(f"文件上传成功: {response.json()}")
            return True
        else:
            print(f"文件上传失败: {response.status_code} - {response.text}")
            return False

# 检索电话号码
def search_phone_number(token):
    url = "http://localhost:8001/rag-query"
    headers = {
        "Authorization": f"Bearer {token}",
        "Content-Type": "application/json"
    }
    data = {
        "query": target_phone,
        "n_results": 5
    }
    response = requests.post(url, headers=headers, json=data)
    if response.status_code == 200:
        results = response.json()
        print(f"API检索结果: {json.dumps(results, ensure_ascii=False, indent=2)}")
        return results
    else:
        print(f"API检索失败: {response.status_code} - {response.text}")
        return None

# 执行API测试
token = get_token()
if token:
    if upload_test_file(token):
        print("等待文件处理完成...")
        time.sleep(5)  # 增加等待时间确保文件被处理
        api_results = search_phone_number(token)
        if api_results and "answer" in api_results:
            print(f"API AI回答: {api_results['answer']}")
            if expected_company in api_results['answer'] and target_phone in api_results['answer']:
                print(f"API测试成功! 正确识别了{target_phone}属于{expected_company}")
            else:
                print(f"API测试失败: 未能识别关联信息")
    else:
        print("文件上传失败，无法进行API测试")
else:
    print("无法获取令牌，无法进行API测试")

print("\n=== 测试完成 ===")